#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
任务队列
按照配置的最大进程数控制任务执行，确保完成一个任务后再进行下一个
"""

import asyncio
import json
from typing import Dict, Any, List, Set
import logging

from app.core.database import SessionLocal
from app.models.task import Task, TaskStatus
from app.services.task_service import get_task_service

logger = logging.getLogger(__name__)


class TaskQueue:
    """任务队列 - 按照配置的最大进程数控制任务执行"""
    
    def __init__(self, ffmpeg_service=None, max_concurrent_tasks=4, poll_interval=2):
        # 使用传入的ffmpeg_service，避免循环导入
        if ffmpeg_service is None:
            from app.services.ffmpeg_service import get_ffmpeg_service
            self.ffmpeg_service = get_ffmpeg_service()
        else:
            self.ffmpeg_service = ffmpeg_service
            
        self.task_service = get_task_service()
        self.running = False
        self.main_task = None
        self.max_concurrent_tasks = max_concurrent_tasks
        self.poll_interval = poll_interval  # 轮询间隔（秒）        
        # 通过running_tasks集合和available_slots计算来控制并发               
        self.running_tasks: Set[str] = set()        
        # 任务处理锁，避免重复处理同一个任务
        self.task_lock = asyncio.Lock()
    
    async def start(self):
        """启动任务队列"""
        if self.running:
            return
            
        self.running = True
        logger.info(f"任务队列已启动，最大并发数: {self.max_concurrent_tasks}")
        
        # 启动主循环任务
        self.main_task = asyncio.create_task(self._main_loop())
    
    async def stop(self):
        """停止任务队列"""
        if not self.running:
            return
            
        self.running = False
        logger.info("任务队列正在停止...")
        
        if self.main_task:
            self.main_task.cancel()
            try:
                await self.main_task
            except asyncio.CancelledError:
                pass
            finally:
                self.main_task = None
        
        # 等待所有正在运行的任务完成
        if self.running_tasks:
            logger.info(f"等待 {len(self.running_tasks)} 个任务完成...")
            while self.running_tasks:
                await asyncio.sleep(1)
        
        logger.info("任务队列已停止")
        
    async def _stop_task(self, task_id: int):
        """停止指定任务（保持向后兼容）"""
        task_id_str = str(task_id)
        logger.info(f"尝试停止任务: {task_id}")
        
        # 通过FFmpeg服务停止任务进程
        success = self.ffmpeg_service.stop_task(task_id_str)
        
        if success:
            # 更新任务状态为已取消
            try:
                await self.task_service.update_task_status(
                    task_id=task_id_str,
                    status=TaskStatus.CANCELLED
                )
                logger.info(f"任务 {task_id} 已停止并更新状态为已取消")
            except Exception as e:
                logger.error(f"更新任务状态失败: {str(e)}")
        else:
            logger.warning(f"任务 {task_id} 停止失败")
        
        # 从运行中任务集合移除
        if task_id_str in self.running_tasks:
            self.running_tasks.remove(task_id_str)
            logger.info(f"任务 {task_id} 已从运行中任务集合移除")
    
    async def _main_loop(self):
        """主循环 - 定期检查待处理任务"""
        while self.running:
            try:
                # 检查并处理待处理任务
                await self._check_and_process_tasks()
                
                # 等待下一次检查
                await asyncio.sleep(self.poll_interval)
                
            except asyncio.CancelledError:
                # 任务被取消，正常退出
                break
            except Exception as e:
                logger.error(f"主循环异常: {str(e)}")
                # 出错时增加等待时间，避免频繁重试
                await asyncio.sleep(max(self.poll_interval, 5))
    
    async def _check_and_process_tasks(self):
        """检查并处理待处理任务"""
        # 计算当前可以启动的新任务数量
        available_slots = self.max_concurrent_tasks - len(self.running_tasks)
        
        if available_slots <= 0:
            # 没有可用槽位，等待
            return
        
        try:
            # 获取待处理任务
            async with self.task_lock:
                # 只获取可用槽位数的任务
                pending_tasks = await self.task_service.get_pending_tasks()                
            if not pending_tasks:
                # 没有待处理任务
                return
            
            # 启动任务处理
            for task in pending_tasks:
                # 使用异步任务处理每个任务，避免阻塞主循环
                print(f"处理任务 {task}")
                asyncio.create_task(self._process_task(task))
                
        except Exception as e:
            logger.error(f"获取待处理任务失败: {str(e)}")
    
    async def _process_task(self, task: Dict[str, Any]) -> None:
        """
        处理单个任务    单独启动一个ffmpeg进程去执行，然后定时监控进程状态  进程不在表示结束或者失败
        """
        task_id = str(task.get("id"))
        task_type = task.get("task_type")        
        # 添加到正在运行的任务集合
        self.running_tasks.add(task_id)
        
        # 从配置中获取命令和输入源信息
        config = task.get("config")
        command = str(config).strip() if config else ""
        input_source = None
        
        # 尝试从命令中提取输入源
        if command and "-i" in command:
            parts = command.split("-i")
            if len(parts) > 1:
                input_source_part = parts[1].strip().split()[0]
                if input_source_part.startswith('"') and input_source_part.endswith('"'):
                    input_source = input_source_part[1:-1]
                elif input_source_part.startswith("'") and input_source_part.endswith("'"):
                    input_source = input_source_part[1:-1]
                else:
                    input_source = input_source_part
        
        try:
            # 记录任务开始时间
            import datetime
            start_time = datetime.datetime.now()
            
            # 验证命令是否包含必要参数
            if not command or "-i" not in command:
                error_msg = f"配置不完整，缺少必要的参数: {command}"
                await self.task_service.update_task_status(
                    task_id=task_id,
                    status=TaskStatus.FAILED,
                    error_message=error_msg
                )
                raise ValueError(error_msg)
            
            # 启动任务
            logger.info(f"启动任务: {task_id} ({task_type}) 命令: {command}")
            
            # 调用start_task，返回结果是字典
            task_result = await self.ffmpeg_service.start_task(
                command=command,
                task_id=task_id,
                timeout=60
            )
            
            # 记录任务启动结果
            print(f"[res]任务 {task_id} 启动结果: {task_result}")
            
            # 检查任务是否启动成功
            if task_result.get("success") != 1:
                error_msg = task_result.get("error", "FFmpeg进程启动失败")
                logger.error(f"任务 {task_id} 启动失败: {error_msg}")
                try:
                    await self.task_service.update_task_status(
                        task_id=task_id,
                        status=TaskStatus.FAILED,
                        error_message=error_msg
                    )
                    logger.info(f"任务 {task_id} 状态已更新为失败: {error_msg}")
                except Exception as e:
                    logger.error(f"更新任务 {task_id} 状态为失败失败: {str(e)}")
                raise ValueError(error_msg)
            
            # 获取任务的进程ID
            pid = task_result.get('pid')
            
            # 记录任务信息
            if pid is not None:
                logger.info(f"任务 {task_id} 进程已启动，PID: {pid}")
                # 更新任务信息
                try:
                    await self.task_service.update_task_info(
                        task_id=task_id,
                        pid=pid,
                        info={'start_time': start_time.isoformat()}
                    )
                except Exception as e:
                    logger.warning(f"更新任务信息失败: {str(e)}")
            
            # 设置输入源访问超时检测
            # 开始监控任务状态
            while True:
                # 检查任务状态
                process_status = self.ffmpeg_service.get_process_status(task_id)
                print(f"[res]任务 {task_id} 进程状态: {process_status}")
                if not process_status or process_status.get('status') != 'running':
                    # 进程已结束或不存在
                    end_time = datetime.datetime.now()
                    duration = (end_time - start_time).total_seconds()
                    
                    # 获取任务退出信息
                    exit_code = process_status.get('return_code', -1) if process_status else -1
                    error_info = process_status.get('error') if process_status else "进程未找到"
                    
                    # 构建资源消耗信息
                    resources = {
                        'cpu_usage': process_status.get('cpu_usage', 0) if process_status else 0,
                        'memory_usage': process_status.get('memory_usage', 0) if process_status else 0,
                        'duration': duration
                    }
                    
                    # 尝试更新任务信息
                    try:                       
                        await self.task_service.update_task_info(
                            task_id=task_id,
                            info={
                                'end_time': end_time.isoformat(),
                                'duration': duration,
                                'resources': resources,
                                'exit_code': exit_code
                            }
                        )
                    except Exception as e:
                        logger.warning(f"更新任务完成信息失败: {str(e)}")
                    
                    # 根据退出码判断任务结果
                    if exit_code == 0:
                        logger.info(f"任务 {task_id} 执行完成，持续时间: {duration:.2f}秒")
                        try:
                            await self.task_service.update_task_status(
                                task_id=task_id,
                                status=TaskStatus.COMPLETED
                            )
                            logger.info(f"任务 {task_id} 状态已更新为完成")
                        except Exception as e:
                            logger.error(f"更新任务 {task_id} 状态为完成失败: {str(e)}")
                    else:
                        # 转码失败，捕获ffmpeg报错信息
                        error_msg = f"转码失败，退出码: {exit_code}, 错误信息: {error_info}"
                        logger.error(error_msg)
                        try:
                            await self.task_service.update_task_status(
                                task_id=task_id,
                                status=TaskStatus.FAILED,
                                error_message=error_msg
                            )
                            logger.info(f"任务 {task_id} 状态已更新为失败: {error_msg}")
                        except Exception as e:
                            logger.error(f"更新任务 {task_id} 状态为失败失败: {str(e)}")
                    break                
                # 休眠一段时间后再次检查
                await asyncio.sleep(60)  # 每60秒检查一次
        except Exception as e:
            error_msg = f"处理任务 {task_id} 异常: {str(e)}"
            logger.error(error_msg)
            
            # 更新任务状态为失败
            try:
                await self.task_service.update_task_status(
                    task_id=task_id,
                    status=TaskStatus.FAILED,
                    error_message=error_msg
                )
                logger.info(f"任务 {task_id} 状态已更新为失败: {error_msg}")
            except Exception as update_err:
                logger.error(f"更新任务状态失败: {str(update_err)}")
        finally:
            # 从运行中的任务集合中移除
            if task_id in self.running_tasks:
                self.running_tasks.remove(task_id)
    

# 全局任务队列实例
_task_queue = None

def get_task_queue():
    """获取任务队列实例"""
    global _task_queue
    if _task_queue is None:
        _task_queue = TaskQueue()
    return _task_queue

async def start_task_queue():
    """启动任务队列"""
    queue = get_task_queue()
    await queue.start()

async def stop_task_queue():
    """停止任务队列"""
    queue = get_task_queue()
    await queue.stop()